"""
DRL模型来自ElegantRL: https://github.com/AI4Finance-Foundation/ElegantRL
"""

# 启用未来版本的特性，确保代码在未来Python版本中的兼容性
from __future__ import annotations

# 导入PyTorch库，用于深度学习相关操作
import torch
# 从ElegantRL的agents模块导入所有内容
from elegantrl.agents import *
# 从ElegantRL的配置模块导入Config类，用于配置训练参数
from elegantrl.train.config import Config
# 从ElegantRL的运行模块导入train_agent函数，用于训练智能体
from elegantrl.train.run import train_agent

# 定义模型字典，将模型名称映射到对应的ElegantRL智能体类
MODELS = {
    "ddpg": AgentDDPG,
    "td3": AgentTD3,
    "sac": AgentSAC,
    "ppo": AgentPPO,
    "a2c": AgentA2C,
}
# 定义离线策略模型列表，包含使用经验回放的模型
OFF_POLICY_MODELS = ["ddpg", "td3", "sac"]
# 定义在线策略模型列表，包含不使用经验回放的模型
ON_POLICY_MODELS = ["ppo"]
# 注释掉的代码，用于从配置中获取各模型的参数
# MODEL_KWARGS = {x: config.__dict__[f"{x.upper()}_PARAMS"] for x in MODELS.keys()}
#
# 注释掉的代码，定义噪声类型字典
# NOISE = {
#     "normal": NormalActionNoise,
#     "ornstein_uhlenbeck": OrnsteinUhlenbeckActionNoise,
# }

class DRLAgent:
    """DRL算法的实现类
    属性
    ----------
        env: gym环境类
            用户自定义的环境类
    方法
    -------
        get_model()
            设置DRL算法模型
        train_model()
            在训练数据集上训练DRL算法，并输出训练好的模型
        DRL_prediction()
            在测试数据集上进行预测并获取结果
    """

    def __init__(self, env, price_array, tech_array, turbulence_array):
        # 初始化环境
        self.env = env
        # 初始化价格数组
        self.price_array = price_array
        # 初始化技术指标数组
        self.tech_array = tech_array
        # 初始化波动率数组
        self.turbulence_array = turbulence_array

    def get_model(self, model_name, model_kwargs):
        # 构建环境配置字典，包含价格、技术指标、波动率等信息，并标记为训练模式
        self.env_config = {
            "price_array": self.price_array,
            "tech_array": self.tech_array,
            "turbulence_array": self.turbulence_array,
            "if_train": True,
        }
        # 保存模型参数
        self.model_kwargs = model_kwargs
        # 从模型参数中获取折扣因子gamma，若未提供则使用默认值0.985
        self.gamma = model_kwargs.get("gamma", 0.985)

        # 获取环境实例
        env = self.env
        # 设置环境数量为1
        env.env_num = 1
        # 根据模型名称从模型字典中获取对应的智能体类
        agent = MODELS[model_name]
        # 检查模型名称是否在支持的模型列表中，若不在则抛出未实现错误
        if model_name not in MODELS:
            raise NotImplementedError("NotImplementedError")

        # 获取股票维度，即股票数量
        stock_dim = self.price_array.shape[1]
        # 计算状态维度
        self.state_dim = 1 + 2 + 3 * stock_dim + self.tech_array.shape[1]
        # 计算动作维度，与股票数量相同
        self.action_dim = stock_dim
        # 构建环境参数字典，包含环境名称、配置、状态维度、动作维度等信息
        self.env_args = {
            "env_name": "StockEnv",
            "config": self.env_config,
            "state_dim": self.state_dim,
            "action_dim": self.action_dim,
            "if_discrete": False,
            "max_step": self.price_array.shape[0] - 1,
        }

        # 创建配置实例，指定智能体类、环境类和环境参数
        model = Config(agent_class=agent, env_class=env, env_args=self.env_args)
        # 判断模型是否为离线策略模型，并设置相应标志
        model.if_off_policy = model_name in OFF_POLICY_MODELS
        # 若提供了模型参数，则对配置进行更新
        if model_kwargs is not None:
            try:
                # 设置训练中断步数，当总步数超过该值时中断训练
                model.break_step = int(
                    2e5
                )  # break training if 'total_step > break_step'
                # 设置多层感知机的中间层维度
                model.net_dims = (
                    128,
                    64,
                )  # the middle layer dimension of MultiLayer Perceptron
                # 设置折扣因子，用于计算未来奖励的折现值
                model.gamma = self.gamma  # discount factor of future rewards
                # 设置水平长度，与最大步数相同
                model.horizon_len = model.max_step
                # 设置重复更新次数，用于使用经验回放缓冲区来减小评论员网络的损失
                model.repeat_times = 16  # repeatedly update network using ReplayBuffer to keep critic's loss small
                # 从模型参数中获取学习率，若未提供则使用默认值1e-4
                model.learning_rate = model_kwargs.get("learning_rate", 1e-4)
                # 设置状态值归一化的tau参数
                model.state_value_tau = 0.1  # the tau of normalize for value and state `std = (1-std)*std + tau*std`
                # 从模型参数中获取评估次数，若未提供则使用默认值2^5
                model.eval_times = model_kwargs.get("eval_times", 2**5)
                # 设置每多少步进行一次评估
                model.eval_per_step = int(2e4)
            except BaseException:
                # 若读取参数失败，则抛出值错误，提示检查模型参数输入
                raise ValueError(
                    "Fail to read arguments, please check 'model_kwargs' input."
                )
        return model

    def train_model(self, model, cwd, total_timesteps=5000):
        # 设置模型的工作目录
        model.cwd = cwd
        # 设置训练的总步数
        model.break_step = total_timesteps
        # 调用train_agent函数进行模型训练
        train_agent(model)

    @staticmethod
    def DRL_prediction(model_name, cwd, net_dimension, environment, env_args):
        # 导入PyTorch库
        import torch

        # 设置GPU ID，>=0表示使用指定ID的GPU，-1表示使用CPU
        gpu_id = 0  # >=0 means GPU ID, -1 means CPU
        # 根据模型名称从模型字典中获取对应的智能体类
        agent_class = MODELS[model_name]
        # 获取股票维度，即股票数量
        stock_dim = env_args["price_array"].shape[1]
        # 计算状态维度
        state_dim = 1 + 2 + 3 * stock_dim + env_args["tech_array"].shape[1]
        # 计算动作维度，与股票数量相同
        action_dim = stock_dim
        # 构建环境参数字典，包含环境数量、名称、状态维度、动作维度等信息
        env_args = {
            "env_num": 1,
            "env_name": "StockEnv",
            "state_dim": state_dim,
            "action_dim": action_dim,
            "if_discrete": False,
            "max_step": env_args["price_array"].shape[0] - 1,
            "config": env_args,
        }

        # 定义演员网络的保存路径
        actor_path = f"{cwd}/act.pth"
        # 设置网络维度
        net_dim = [net_dimension]

        """初始化"""
        # 获取环境实例
        env = environment
        # 获取环境类
        env_class = env
        # 创建配置实例，指定智能体类、环境类和环境参数
        args = Config(agent_class=agent_class, env_class=env_class, env_args=env_args)
        # 设置配置的工作目录
        args.cwd = cwd
        # 初始化演员网络
        act = agent_class(
            net_dim, env.state_dim, env.action_dim, gpu_id=gpu_id, args=args
        ).act
        # 初始化参数字典
        parameters_dict = {}
        # 从文件中加载演员网络的参数
        act = torch.load(actor_path)
        # 遍历演员网络的参数，将其转换为张量并存入参数字典
        for name, param in act.named_parameters():
            parameters_dict[name] = torch.tensor(param.detach().cpu().numpy())

        # 将参数字典中的参数加载到演员网络中
        act.load_state_dict(parameters_dict)

        # 判断动作是否为离散动作
        if_discrete = env.if_discrete
        # 获取网络参数所在的设备
        device = next(act.parameters()).device
        # 重置环境，获取初始状态
        state = env.reset()
        # 初始化每一步的累计回报列表
        episode_returns = []  # the cumulative_return / initial_account
        # 初始化每一步的总资产列表，初始值为初始总资产
        episode_total_assets = [env.initial_total_asset]
        # 获取环境的最大步数
        max_step = env.max_step
        # 循环执行最大步数
        for steps in range(max_step):
            # 将当前状态转换为张量，并添加一个批次维度
            s_tensor = torch.as_tensor(
                state, dtype=torch.float32, device=device
            ).unsqueeze(0)
            # 根据动作是否离散，选择不同的方式获取动作张量
            a_tensor = act(s_tensor).argmax(dim=1) if if_discrete else act(s_tensor)
            # 将动作张量转换为numpy数组
            action = (
                a_tensor.detach().cpu().numpy()[0]
            )  # not need detach(), because using torch.no_grad() outside
            # 执行动作，获取下一个状态、奖励、是否结束等信息
            state, reward, done, _ = env.step(action)
            # 计算当前总资产
            total_asset = env.amount + (env.price_ary[env.day] * env.stocks).sum()
            # 将当前总资产添加到总资产列表中
            episode_total_assets.append(total_asset)
            # 计算当前累计回报
            episode_return = total_asset / env.initial_total_asset
            # 将当前累计回报添加到累计回报列表中
            episode_returns.append(episode_return)
            # 若游戏结束，则跳出循环
            if done:
                break
        # 打印测试结束信息
        print("Test Finished!")
        # 打印最后一步的累计回报
        print("episode_retuen", episode_return)
        # 返回每一步的总资产列表
        return episode_total_assets
